package cn.wmelon.shell.executor;


import ch.ethz.ssh2.ChannelCondition;
import ch.ethz.ssh2.Connection;
import ch.ethz.ssh2.Session;
import ch.ethz.ssh2.StreamGobbler;
import cn.wmelon.shell.authentication.ShellAuthenticationEntity;
import cn.wmelon.shell.command.Command;
import cn.wmelon.shell.entity.ShellResultEntity;

import java.io.*;
import java.nio.charset.StandardCharsets;

/**
 * @author watermelon
 * @time 2020/2/18
 */
public class RemoteShellExecutor {

    /**
     * 超时时间 毫秒 （可在 ShellAuthenticationEntity 中设置）
     * {@link ShellAuthenticationEntity#setTimeout(int)} ()}
     */
    private static final int TIME_OUT = 360000;


    /**
     * 执行脚本
     * 如果遇到环境变量不完全的问题
     * 可使用{@link #execPTY(String, ShellAuthenticationEntity)} {@link #execPTY(Command, ShellAuthenticationEntity)}
     *
     * @param shell
     * @return
     * @throws Exception
     */
    public static ShellResultEntity exec(String shell, ShellAuthenticationEntity se) {
        ShellResultEntity shellResultEntity = new ShellResultEntity();
        Connection conn = new Connection(se.getIp(), se.getPort());
        try {
            conn.connect();
            if (conn.authenticateWithPassword(se.getUsername(), se.getPassword())) {
                Session session = conn.openSession();
                session.execCommand(shell);
                session.waitForCondition(ChannelCondition.EXIT_STATUS, se.getTimeout() == 0 ? TIME_OUT : se.getTimeout());
                shellResultEntity.setCode(session.getExitStatus());
                shellResultEntity.setMsg(processStdout(session.getStdout()));
                shellResultEntity.setErrorMsg(processStdout(session.getStderr()));
            } else {
                throw new RuntimeException("登录远程机器失败" + se.getIp());
            }
        } catch (Exception e) {
            throw new RuntimeException("执行脚本 " + shell + " 失败");
        } finally {
            conn.close();
        }
        return shellResultEntity;
    }


    /**
     * 执行脚本
     * 如果遇到环境变量不完全的问题
     * 可使用{@link #execPTY(String, ShellAuthenticationEntity)} {@link #execPTY(Command, ShellAuthenticationEntity)}
     *
     * @param command
     * @return
     * @throws Exception
     */
    public static ShellResultEntity exec(Command command, ShellAuthenticationEntity se) {
        return exec(command.getCommandString(), se);
    }

    /**
     * 执行脚本
     * （todo 可能有问题）
     *
     * @param shell
     * @return
     * @throws Exception
     */
    public static ShellResultEntity execPTY(String shell, ShellAuthenticationEntity se) {
        ShellResultEntity shellResultEntity = new ShellResultEntity();
        Connection conn = new Connection(se.getIp(), se.getPort());
        try {
            conn.connect();
            if (conn.authenticateWithPassword(se.getUsername(), se.getPassword())) {
                Session session = conn.openSession();
                // 建立虚拟终端
                session.requestPTY("bash");
                // 打开一个Shell
                session.startShell();
                // 准备输入命令
                PrintWriter out = new PrintWriter(session.getStdin());
                // 输入待执行命令
                out.println(shell);
                out.println("exit");
                // 6. 关闭输入流
                out.close();
                // 7. 等待，除非1.连接关闭；2.输出数据传送完毕；3.进程状态为退出；4.超时
                session.waitForCondition(ChannelCondition.CLOSED | ChannelCondition.EOF | ChannelCondition.EXIT_STATUS, TIME_OUT);

                shellResultEntity.setCode(session.getExitStatus());
                shellResultEntity.setMsg(processStdout(session.getStdout()));
                shellResultEntity.setErrorMsg(processStdout(session.getStderr()));
            } else {
                throw new RuntimeException("登录远程机器失败" + se.getIp());
            }
        } catch (Exception e) {
            throw new RuntimeException("执行脚本 " + shell + " 失败");
        } finally {
            conn.close();
        }
        return shellResultEntity;
    }


    /**
     * 执行脚本
     *
     * @param shell
     * @return
     * @throws Exception
     */
    public static ShellResultEntity execPTY(Command shell, ShellAuthenticationEntity se) {
        return execPTY(shell.getCommandString(), se);
    }

    /**
     * 解析脚本执行返回的结果集
     *
     * @param in 输入流对象
     * @return 以纯文本的格式返回
     */
    private static String processStdout(InputStream in) {
        InputStream stdout = new StreamGobbler(in);
        StringBuffer buffer = new StringBuffer();
        ;
        try {
            BufferedReader br = new BufferedReader(new InputStreamReader(stdout, StandardCharsets.UTF_8));
            String line = null;
            while ((line = br.readLine()) != null) {
                buffer.append(line + "\n");
                System.out.println(line);
            }
            br.close();
        } catch (UnsupportedEncodingException e) {
            System.out.println("解析脚本出错：" + e.getMessage());
            e.printStackTrace();
        } catch (IOException e) {
            System.out.println("解析脚本出错：" + e.getMessage());
            e.printStackTrace();
        }
        return buffer.toString();
    }

}
